from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QSplitter,
    QGraphicsView, QGraphicsScene, QToolBar,
     QComboBox, QPushButton, QLabel,
    QTableWidget, QTextEdit, QListWidget, QListWidgetItem,QGraphicsItem
)
from PyQt6.QtGui import QIcon, QPen, QBrush, QColor, QFont,QAction
from PyQt6.QtCore import Qt, QPointF, QRectF
from core.connection import ConnectionManager
from core.query_engine import QueryEngine
from utils.logger import get_logger
from PyQt6.QtWidgets import QGraphicsItemGroup, QGraphicsRectItem, QGraphicsTextItem
logger = get_logger(__name__)

class TableNode(QGraphicsItem):
    def __init__(self, table_name, columns, x=0, y=0):
        super().__init__()
        self.table_name = table_name
        self.columns = columns
        self.width = 180
        self.height = 40 + len(columns) * 25
        self.setPos(x, y)
        self.setFlag(QGraphicsItem.GraphicsItemFlag.ItemIsMovable)
        self.setFlag(QGraphicsItem.GraphicsItemFlag.ItemIsSelectable)
        self.setAcceptHoverEvents(True)

    def boundingRect(self):
        return QRectF(0, 0, self.width, self.height)

    def paint(self, painter, option, widget):
        # 绘制表头
        painter.setBrush(QBrush(QColor(200, 220, 255)))
        painter.setPen(QPen(Qt.GlobalColor.black, 1))
        painter.drawRect(0, 0, self.width, 30)
        painter.setFont(QFont("Arial", 10, QFont.Weight.Bold))
        painter.drawText(10, 20, self.table_name)

        # 绘制列
        painter.setBrush(QBrush(QColor(240, 240, 240)))
        for i, col in enumerate(self.columns):
            y_pos = 30 + i * 25
            painter.drawRect(0, y_pos, self.width, 25)
            painter.setPen(QPen(Qt.GlobalColor.black))
            painter.setFont(QFont("Arial", 9))
            painter.drawText(10, y_pos + 18, col)

        # 如果被选中，绘制边框
        if self.isSelected():
            painter.setPen(QPen(Qt.GlobalColor.blue, 2))
            painter.drawRect(0, 0, self.width, self.height)

class JoinLine(QGraphicsItem):
    def __init__(self, from_node, from_column, to_node, to_column):
        super().__init__()
        self.from_node = from_node
        self.from_column = from_column
        self.to_node = to_node
        self.to_column = to_column
        self.setZValue(-1)  # 确保线在节点下方

    def boundingRect(self):
        return QRectF(self.from_node.scenePos(), self.to_node.scenePos()).normalized()

    def paint(self, painter, option, widget):
        from_point = self.from_node.scenePos() + QPointF(
            self.from_node.boundingRect().width() / 2,
            30 + self.from_column * 25 + 12
        )
        to_point = self.to_node.scenePos() + QPointF(
            self.to_node.boundingRect().width() / 2,
            30 + self.to_column * 25 + 12
        )

        painter.setPen(QPen(Qt.GlobalColor.darkGray, 2))
        painter.drawLine(from_point, to_point)

class VisualQueryBuilder(QWidget):
    def __init__(self, connection=None, parent=None):
        super().__init__(parent)
        self.connection = connection
        self.scene = QGraphicsScene()
        self.nodes = []
        self.lines = []
        self.selected_column = None

        # 创建布局
        layout = QVBoxLayout()
        layout.setContentsMargins(5, 5, 5, 5)
        self.setLayout(layout)

        # 创建工具栏
        self.create_toolbar()
        layout.addWidget(self.toolbar)

        # 创建主分割器
        main_splitter = QSplitter(Qt.Orientation.Horizontal)
        layout.addWidget(main_splitter)

        # 左侧：数据库对象树
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)

        # 可用表列表
        left_layout.addWidget(QLabel("可用表:"))
        self.table_list = QListWidget()
        self.table_list.itemDoubleClicked.connect(self.add_table_to_canvas)
        left_layout.addWidget(self.table_list)

        # 字段列表
        left_layout.addWidget(QLabel("表字段:"))
        self.column_list = QListWidget()
        self.column_list.itemDoubleClicked.connect(self.select_column_for_join)
        left_layout.addWidget(self.column_list)

        # 连接类型选择
        join_type_layout = QHBoxLayout()
        join_type_layout.addWidget(QLabel("连接类型:"))
        self.join_type_combo = QComboBox()
        self.join_type_combo.addItems(["INNER JOIN", "LEFT JOIN", "RIGHT JOIN", "FULL JOIN"])
        join_type_layout.addWidget(self.join_type_combo)
        left_layout.addLayout(join_type_layout)

        # 添加连接按钮
        self.join_button = QPushButton("添加连接")
        self.join_button.setEnabled(False)
        self.join_button.clicked.connect(self.add_join)
        left_layout.addWidget(self.join_button)

        # 生成SQL按钮
        generate_sql_button = QPushButton("生成SQL")
        generate_sql_button.clicked.connect(self.generate_sql)
        left_layout.addWidget(generate_sql_button)

        main_splitter.addWidget(left_panel)

        # 中间：画布区域
        canvas_panel = QWidget()
        canvas_layout = QVBoxLayout(canvas_panel)
        canvas_layout.addWidget(QLabel("查询设计画布:"))

        self.view = QGraphicsView(self.scene)
        self.view.setRenderHint(QPainter.RenderHint.Antialiasing)
        self.view.setSceneRect(-1000, -1000, 2000, 2000)
        canvas_layout.addWidget(self.view)

        main_splitter.addWidget(canvas_panel)

        # 右侧：SQL预览和结果
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)

        # SQL预览
        right_layout.addWidget(QLabel("SQL预览:"))
        self.sql_preview = QTextEdit()
        self.sql_preview.setReadOnly(True)
        self.sql_preview.setFont(QFont("Consolas", 10))
        right_layout.addWidget(self.sql_preview)

        # 执行按钮
        execute_button = QPushButton("执行查询")
        execute_button.clicked.connect(self.execute_query)
        right_layout.addWidget(execute_button)

        # 结果表格
        right_layout.addWidget(QLabel("查询结果:"))
        self.result_table = QTableWidget()
        right_layout.addWidget(self.result_table)

        main_splitter.addWidget(right_panel)

        # 设置分割比例
        main_splitter.setSizes([200, 600, 400])

        # 如果提供了连接，加载表
        if connection:
            self.load_tables()

    def create_toolbar(self):
        """创建工具栏"""
        self.toolbar = QToolBar("可视化查询工具栏")
        self.toolbar.setIconSize(QSize(20, 20))

        # 添加表按钮
        add_table_action = QAction(QIcon("resources/icons/add_table.png"), "添加表", self)
        add_table_action.triggered.connect(self.add_selected_table)
        self.toolbar.addAction(add_table_action)

        # 清除按钮
        clear_action = QAction(QIcon("resources/icons/clear.png"), "清除画布", self)
        clear_action.triggered.connect(self.clear_canvas)
        self.toolbar.addAction(clear_action)

        # 删除按钮
        delete_action = QAction(QIcon("resources/icons/delete.png"), "删除选中", self)
        delete_action.triggered.connect(self.delete_selected)
        self.toolbar.addAction(delete_action)

    def load_tables(self):
        """加载数据库中的表"""
        if not self.connection:
            return

        try:
            # 根据数据库类型获取表列表
            db_type = self.connection["type"]
            if db_type == "mysql":
                import pymysql
                conn = pymysql.connect(
                    host=self.connection["host"],
                    user=self.connection["user"],
                    password=self.connection["password"],
                    database=self.connection.get("database", ""),
                    port=self.connection.get("port", 3306)
                )
                cursor = conn.cursor()
                cursor.execute("SHOW TABLES")
                tables = [table[0] for table in cursor.fetchall()]
                cursor.close()
                conn.close()
            elif db_type == "postgres":
                import psycopg2
                conn = psycopg2.connect(
                    host=self.connection["host"],
                    user=self.connection["user"],
                    password=self.connection["password"],
                    dbname=self.connection.get("database", ""),
                    port=self.connection.get("port", 5432)
                )
                cursor = conn.cursor()
                cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='public'")
                tables = [table[0] for table in cursor.fetchall()]
                cursor.close()
                conn.close()
            else:
                tables = []

            # 更新表列表
            self.table_list.clear()
            for table in tables:
                self.table_list.addItem(table)
        except Exception as e:
            logger.error(f"加载表列表错误: {str(e)}")

    def add_selected_table(self):
        """添加选中的表到画布"""
        selected_items = self.table_list.selectedItems()
        if selected_items:
            table_name = selected_items[0].text()
            self.add_table_to_canvas(table_name)

    def add_table_to_canvas(self, item):
        """添加表到画布"""
        if isinstance(item, QListWidgetItem):
            table_name = item.text()
        else:
            table_name = item

        # 获取表列信息
        columns = self.get_table_columns(table_name)

        # 创建表节点
        x = len(self.nodes) * 250
        node = TableNode(table_name, columns, x, 100)
        self.scene.addItem(node)
        self.nodes.append(node)

    def get_table_columns(self, table_name):
        """获取表的列信息"""
        if not self.connection:
            return []

        try:
            db_type = self.connection["type"]
            if db_type == "mysql":
                import pymysql
                conn = pymysql.connect(
                    host=self.connection["host"],
                    user=self.connection["user"],
                    password=self.connection["password"],
                    database=self.connection.get("database", ""),
                    port=self.connection.get("port", 3306)
                )
                cursor = conn.cursor()
                cursor.execute(f"DESCRIBE {table_name}")
                columns = [col[0] for col in cursor.fetchall()]
                cursor.close()
                conn.close()
                return columns
            elif db_type == "postgres":
                import psycopg2
                conn = psycopg2.connect(
                    host=self.connection["host"],
                    user=self.connection["user"],
                    password=self.connection["password"],
                    dbname=self.connection.get("database", ""),
                    port=self.connection.get("port", 5432)
                )
                cursor = conn.cursor()
                cursor.execute(f"""
                    SELECT column_name 
                    FROM information_schema.columns 
                    WHERE table_name = '{table_name}'
                """)
                columns = [col[0] for col in cursor.fetchall()]
                cursor.close()
                conn.close()
                return columns
            else:
                return []
        except Exception as e:
            logger.error(f"获取表列错误: {str(e)}")
            return []

    def select_column_for_join(self, item):
        """选择用于连接的列"""
        if not item:
            return

        # 获取选中的表
        selected_table_items = self.table_list.selectedItems()
        if not selected_table_items:
            return

        table_name = selected_table_items[0].text()
        column_name = item.text()

        # 检查是否有选中的节点
        selected_nodes = [node for node in self.nodes if node.isSelected()]
        if not selected_nodes:
            # 如果没有选中的节点，保存选择的列用于后续连接
            self.selected_column = (table_name, column_name)
            self.join_button.setEnabled(True)
            return

        # 如果已经有选中的节点，直接创建连接
        for node in selected_nodes:
            if node.table_name != table_name:
                # 找到列索引
                column_idx = node.columns.index(column_name) if column_name in node.columns else 0
                self.add_join_between(self.selected_column[0], self.selected_column[1], node.table_name, column_name)
                self.selected_column = None
                self.join_button.setEnabled(False)
                break

    def add_join(self):
        """添加连接"""
        if not self.selected_column:
            return

        # 获取选中的节点
        selected_nodes = [node for node in self.nodes if node.isSelected()]
        if not selected_nodes:
            return

        for node in selected_nodes:
            if node.table_name != self.selected_column[0]:
                # 找到列索引
                column_idx = node.columns.index(self.selected_column[1]) if self.selected_column[1] in node.columns else 0
                self.add_join_between(self.selected_column[0], self.selected_column[1], node.table_name, self.selected_column[1])
                self.selected_column = None
                self.join_button.setEnabled(False)
                break

    def add_join_between(self, from_table, from_column, to_table, to_column):
        """在两个表之间添加连接"""
        from_node = None
        to_node = None

        # 查找节点
        for node in self.nodes:
            if node.table_name == from_table:
                from_node = node
            elif node.table_name == to_table:
                to_node = node

        if not from_node or not to_node:
            return

        # 查找列索引
        try:
            from_col_idx = from_node.columns.index(from_column)
            to_col_idx = to_node.columns.index(to_column)
        except ValueError:
            from_col_idx = 0
            to_col_idx = 0

        # 创建连接线
        line = JoinLine(from_node, from_col_idx, to_node, to_col_idx)
        self.scene.addItem(line)
        self.lines.append(line)

    def clear_canvas(self):
        """清除画布"""
        self.scene.clear()
        self.nodes = []
        self.lines = []
        self.selected_column = None
        self.join_button.setEnabled(False)

    def delete_selected(self):
        """删除选中的元素"""
        # 删除选中的节点
        for node in self.nodes[:]:
            if node.isSelected():
                # 删除与节点相关的连接线
                for line in self.lines[:]:
                    if line.from_node == node or line.to_node == node:
                        self.scene.removeItem(line)
                        self.lines.remove(line)

                self.scene.removeItem(node)
                self.nodes.remove(node)

        # 删除选中的连接线
        for line in self.lines[:]:
            if line.isSelected():
                self.scene.removeItem(line)
                self.lines.remove(line)

    def generate_sql(self):
        """生成SQL查询"""
        if not self.nodes:
            self.sql_preview.setText("请添加表到画布")
            return

        # 构建SELECT子句
        select_clause = "SELECT "
        columns = []
        for node in self.nodes:
            for col in node.columns:
                columns.append(f"{node.table_name}.{col}")
        select_clause += ", ".join(columns) + "\n"

        # 构建FROM子句
        from_clause = "FROM " + self.nodes[0].table_name + "\n"

        # 构建JOIN子句
        join_clauses = []
        for line in self.lines:
            join_type = self.join_type_combo.currentText()
            join_clauses.append(
                f"{join_type} {line.to_node.table_name} "
                f"ON {line.from_node.table_name}.{line.from_node.columns[line.from_column]} = "
                f"{line.to_node.table_name}.{line.to_node.columns[line.to_column]}"
            )

        # 组合SQL
        sql = select_clause + from_clause + "\n".join(join_clauses) + ";"
        self.sql_preview.setText(sql)

    def execute_query(self):
        """执行查询"""
        sql = self.sql_preview.toPlainText()
        if not sql or not self.connection:
            return

        # 执行查询
        query_engine = QueryEngine()
        try:
            result = query_engine.execute(self.connection, sql)

            if result["type"] == "query":
                # 显示结果
                self.display_results(result["data"], result["columns"])
            else:
                self.sql_preview.append(f"\n\n{result.get('message', '命令执行成功')}")
        except Exception as e:
            self.sql_preview.append(f"\n\n执行错误: {str(e)}")

    def display_results(self, data, columns):
        """在表格中显示结果"""
        self.result_table.clear()
        self.result_table.setRowCount(len(data))
        self.result_table.setColumnCount(len(columns))
        self.result_table.setHorizontalHeaderLabels(columns)

        for row_idx, row in enumerate(data):
            for col_idx, value in enumerate(row):
                self.result_table.setItem(row_idx, col_idx, QTableWidgetItem(str(value)))